# -*- coding: UTF-8 -*-

from pymongo import MongoClient
from WindPy import *

import time
import datetime
import math



class DailyClose:
    def __init__(self, startDateStr, endDateStr = time.strftime('%Y-%m-%d'), fillType='F', DBIP='localhost', DBPort=27017, DBName = 'stocks', txtFileName = 'close.txt', colStocksName = 'daily_close', colIndicesName = 'indices_close', colSysData = 'sys_data'):
        self.startDateStr = startDateStr
        self.endDateStr = endDateStr
        self.fillType = fillType
        self.DBIP = DBIP
        self.DBPort = DBPort
        self.DBName = DBName
        self.txtFileName = txtFileName
        self.transDateList = []
        self.stockIndexIdList = ['000001.SH', '399001.SZ', '399006.SZ', '000300.SH', '000016.SH', '000905.SH', '399005.SZ', '881001.WI']
        self.colStocksName = colStocksName
        self.colIndicesName = colIndicesName
        self.colSysData = colSysData
        self.stockExtractGrpNum = 8
        w.start()



    def writeBaseStocksFromWindToMongoDB(self, endDateStr):

        ###############################################################################################################
        # format of tje default 'startDateStr' is '2012-01-09'
        # format of tje default 'endDateStr' is '2017-05-08'
        # fillType = 'F': 前复权 ，fillType = 'B': 后复权
        # In order to make data extraction from WindDB faster, it's better to extract all stocks in one Wind retrieval,
        # however too large wind retrieval result causes error, therefore we split stock in to several group,
        # the variable stockExtractGrpNum indicates the group number, the default value is 5 in our case
        ###############################################################################################################

        startDateStr = self.startDateStr
        fillType = self.fillType
        DBIP = self.DBIP
        DBPort = self.DBPort
        DBName = self.DBName
        colName = self.colStocksName

        stockExtractGrpNum = self.stockExtractGrpNum

        resStockID = w.wset('sectorconstituent', 'date=' + endDateStr + ';sectorid=a001010100000000;field=wind_code')

        stockIDdata = resStockID.Data[0]

        totalStockIdNum = len(stockIDdata)

        ###############################################################################################################
        #this shows stock number to retrieve in one group
        ###############################################################################################################

        splitStocksNum = math.floor(totalStockIdNum / stockExtractGrpNum)

        client = MongoClient(DBIP, DBPort)

        db = client[DBName]

        collection_dailyclose = db[colName]


        for i in range(0, stockExtractGrpNum):
            stockIDOneLoopList = []
            for j in range(0, splitStocksNum):
                stockIDOneLoopList.append(i*splitStocksNum + j)
            oneStockSearchStr = ''
            for aStockID in stockIDOneLoopList:
                oneStockSearchStr = oneStockSearchStr + str(stockIDdata[aStockID]) + ', '
            resDailyClose = w.wsd(oneStockSearchStr, 'close', startDateStr, endDateStr, 'Fill=Previous;PriceAdj=' + fillType)

            batchStockNum = len(resDailyClose.Data)
            closeDataTime = resDailyClose.Times
            self.transDateList = closeDataTime

            for k in range(0, batchStockNum):
                aData = []
                aDateDict = {}
                aDateDict['StockID'] = stockIDdata[i * splitStocksNum + k]
                totalTransDaysNum = len(closeDataTime)
                for l in range(0, totalTransDaysNum):
                    aDateDict[closeDataTime[l].strftime('%Y%m%d')] = resDailyClose.Data[k][l]
                aData.append(aDateDict)
                collection_dailyclose.insert(aData)
            print('Finished writing stocks of:' + oneStockSearchStr)

        ###############################################################################################################
        # After the looped retrieval, there could be still some stocks not retrieved, i.e., looped retrieved stocks may
        # not be the full set of available stocks, therefore we should locate the unretrieved stocks
        # by using mod function, we can check if remaining stocks is the full set of the available stock set
        # if not, we locate the remains
        ###############################################################################################################
        if (totalStockIdNum % stockExtractGrpNum) != 0:
            stockIDRemainsList = []
            for m in range((stockExtractGrpNum-1)*splitStocksNum + splitStocksNum, totalStockIdNum):
                stockIDRemainsList.append(m)
            oneStockSearchStr = ''

            for aStockID in stockIDRemainsList:
                oneStockSearchStr = oneStockSearchStr + str(stockIDdata[aStockID]) + ', '
            resDailyClose = w.wsd(oneStockSearchStr, 'close', startDateStr, endDateStr, 'Fill=Previous;PriceAdj=' + fillType)
            batchStockNum = len(resDailyClose.Data)
            closeDataTime = resDailyClose.Times

            for k in range(0, batchStockNum):
                aData = []
                aDateDict = {}
                #aDateDict['StockID'] = stockIDRemainsList[k]
                aDateDict['StockID'] = stockIDdata[(stockExtractGrpNum-1)*splitStocksNum + splitStocksNum + k]
                for l in range(0, totalTransDaysNum):
                    aDateDict[closeDataTime[l].strftime('%Y%m%d')] = resDailyClose.Data[k][l]
                aData.append(aDateDict)
                collection_dailyclose.insert(aData)
            print('Finished writing remaining stocks of:' + oneStockSearchStr)
        self.updateDataUpdateDate('last_stock_daily_close_update_date', ' stock close ', currentDateStr = endDateStr)




    ###############################
    ###Update Stocks Close Data
    ##############################
    def updateStocksFromWindToMongoDB(self, startDateStr, endDateStr = time.strftime('%Y-%m-%d')):

        fillType = self.fillType
        DBIP = self.DBIP
        DBPort = self.DBPort
        DBName = self.DBName
        colName = self.colStocksName

        resStockID = w.wset('sectorconstituent', 'date=' + endDateStr + ';sectorid=a001010100000000;field=wind_code')

        stockIDdata = resStockID.Data[0]

        totalStockIdNum = len(stockIDdata)

        stockIndexSearchStr = ''

        for n in range(0, totalStockIdNum):
            stockIndexSearchStr = stockIndexSearchStr + stockIDdata[n] + ', '
        resDailyClose = w.wsd(stockIndexSearchStr, 'close', startDateStr, endDateStr, 'Fill=Previous;PriceAdj=' + fillType)


        closeDataTime = resDailyClose.Times
        if (str(resDailyClose.Data[0][0]).find('No data') != -1):
            print('Currently it is NOT in a transacation day and stock data is up to date, no more data to be updated by now...')
            return()

        totalTransDaysNum = len(closeDataTime)


        client = MongoClient(DBIP, DBPort)
        db = client[DBName]
        collection_dailyclose = db[colName]
        #collection_dailyclose.update({'StockID': '888888.HZ'}, {'$set': {'2017-04-05': '3333'}}, upsert=True)
        dailyCloseIdList = resDailyClose.Data
        if (totalTransDaysNum >1 ):
            for n in range(0, totalStockIdNum):
                aStockIdStr = str(stockIDdata[n])
                for l in range(0, totalTransDaysNum):
                    aDateStr = closeDataTime[l].strftime('%Y%m%d')
                    aCloseStr = str(dailyCloseIdList[n][l])
                    collection_dailyclose.update({'StockID': aStockIdStr}, {'$set': {aDateStr: aCloseStr}}, upsert=True)
                    print('finished updating stock:' + str(stockIDdata[n]) + ', on the day of ' +  aDateStr)
        else:
            for n in range(0, totalStockIdNum):
                aStockIdStr = str(stockIDdata[n])
                for l in range(0, totalTransDaysNum):
                    aDateStr = closeDataTime[l].strftime('%Y%m%d')
                    aCloseStr = str(dailyCloseIdList[l][n])
                    collection_dailyclose.update({'StockID': aStockIdStr}, {'$set': {aDateStr: aCloseStr}}, upsert=True)
                    print('finished updating stock:' + str(stockIDdata[n]) + ', on the day of ' +  aDateStr)
        print('Finished updating all stocks of:' + stockIndexSearchStr)
        self.updateDataUpdateDate('last_stock_daily_close_update_date', ' stock close ')






    def writeBaseIndicesFromWindToMongoDB(self, endDateStr):
        ###############################################################################################################
        #now we add stock index data into MongoDB
        ###############################################################################################################
        startDateStr = self.startDateStr
        fillType = self.fillType
        DBIP = self.DBIP
        DBPort = self.DBPort
        DBName = self.DBName
        ColName = self.colIndicesName

       # closeDataTime = self.transDateList


        stockIndexIdList = self.stockIndexIdList
        stockIndexSearchStr = ''
        stockIndexNum = len(stockIndexIdList)

        for n in range(0, stockIndexNum):
            stockIndexSearchStr = stockIndexSearchStr + stockIndexIdList[n] + ', '

        resDailyClose = w.wsd(stockIndexSearchStr, 'close', startDateStr, endDateStr, 'Fill=Previous;PriceAdj=' + fillType)
        closeDataTime = resDailyClose.Times
        totalTransDaysNum = len(closeDataTime)

        client = MongoClient(DBIP, DBPort)
        db = client[DBName]
        collection_indices_dailyclose = db[ColName]
        for n in range(0, stockIndexNum):
            aData = []
            aDateDict = {}
            aDateDict['StockID'] = stockIndexIdList[n]
            for l in range(0, totalTransDaysNum):
                aDateDict[closeDataTime[l].strftime('%Y%m%d')] = resDailyClose.Data[n][l]
            aData.append(aDateDict)
            collection_indices_dailyclose.insert(aData)
        print('Finished writing stock indices of:' + stockIndexSearchStr)
        self.updateDataUpdateDate('last_index_daily_close_update_date', ' index close ', currentDateStr = endDateStr)






    ###############################
    ###Update Indices Close Data
    ##############################
    def updateIndicesFromWindToMongoDB(self, startDateStr, endDateStr = time.strftime('%Y-%m-%d'), collectionName = 'indices_close'):

        fillType = self.fillType
        DBIP = self.DBIP
        DBPort = self.DBPort
        DBName = self.DBName
        ColName = self.colIndicesName

        stockIndexIdList = self.stockIndexIdList

        stockIndexSearchStr = ''
        stockIndexNum = len(stockIndexIdList)
        for n in range(0, stockIndexNum):
            stockIndexSearchStr = stockIndexSearchStr + stockIndexIdList[n] + ', '
        resDailyClose = w.wsd(stockIndexSearchStr, 'close', startDateStr, endDateStr, 'Fill=Previous;PriceAdj=' + fillType)

        if (str(resDailyClose.Data[0][0]).find('No data') != -1):
            print('Currently it is NOT in a transacation day and index data is up to date, no more data to be updated by now...')
            return()

        closeDataTime = resDailyClose.Times

        totalTransDaysNum = len(closeDataTime)


        client = MongoClient(DBIP, DBPort)
        db = client[DBName]
        collection_indices_dailyclose = db[ColName]
        if (totalTransDaysNum >1 ):
            for n in range(0, stockIndexNum):
                aStockIdStr = str(stockIndexIdList[n])
                for l in range(0, totalTransDaysNum):

                    aDateStr = closeDataTime[l].strftime('%Y%m%d')
                    aCloseStr = str(resDailyClose.Data[n][l])
                    collection_indices_dailyclose.update({'StockID': aStockIdStr}, {'$set': {aDateStr: aCloseStr}}, upsert=True)
                    print('finsihed upating for one transaction day for: ' + str(stockIndexIdList[n]) + ', on the day of ' +  aDateStr)
        else:
            for n in range(0, stockIndexNum):
                aStockIdStr = str(stockIndexIdList[n])
                for l in range(0, totalTransDaysNum):
                    aDateStr = closeDataTime[l].strftime('%Y%m%d')
                    aCloseStr = str(resDailyClose.Data[l][n])
                    collection_indices_dailyclose.update({'StockID': aStockIdStr}, {'$set': {aDateStr: aCloseStr}}, upsert=True)
                    print('finsihed upating for one transaction day for: ' + str(stockIndexIdList[n]) + ', on the day of ' + aDateStr)
        print('Finished updating indices of:' + stockIndexSearchStr)
        self.updateDataUpdateDate('last_index_daily_close_update_date', ' index close ')
        w.stop()



    def writeFromMongoDBtoTxt(self, startDateStr):
        ###############################################################################################################
        # format of the default 'startDateStr' is '2012-01-09'
        # format of the default 'endDateStr' is '2017-05-08'
        # fillType = 'F': 前复权 ，fillType = 'B': 后复权
        ###############################################################################################################

        DBIP = self.DBIP
        DBPort = self.DBPort
        DBName = self.DBName

        txtFileName = self.txtFileName
        colStockName = self.colStocksName
        colIndicesName = self.colIndicesName

        client = MongoClient(DBIP, DBPort)
        db = client[DBName]
        collection_dailyclose = db[colStockName]
        collection_indicesclose = db[colIndicesName]

        dateList = []
        totalStockCloseList = []

        dateListRes = collection_dailyclose.find({'StockID': '000001.SZ'})
        dateListDict = dateListRes[0]

        locateDateStrList =[]

        matchDateStr = datetime.datetime.strptime(startDateStr, '%Y-%m-%d').strftime('%Y%m')

        for (k, v) in dateListDict.items():
            try:
                dateStr = datetime.datetime.strptime(k, '%Y%m%d').strftime('%Y%m%d')
            except:
                dateStr = k

            try:
                locateDateStr = datetime.datetime.strptime(k, '%Y%m%d').strftime('%Y%m')
            except:
                locateDateStr = k

            dateList.append(dateStr)
            locateDateStrList.append(locateDateStr)

        dateList = dateList[2:]
        locateDateStrList = locateDateStrList[2:]

        dateNdx = locateDateStrList.index(matchDateStr)
        dateList = dateList[dateNdx:]
        dateList.insert(0, 'close')
        totalStockCloseList.append(dateList)
        totalSelectedDateNum = len(dateList)

        #######################################
        # Retrieve stock close data
        #######################################
        for stockIdRes in collection_dailyclose.find({}, {'StockID': 1}):
            stockIdStr = stockIdRes['StockID']
            for stockDailyCloseResDict in collection_dailyclose.find({'StockID': stockIdStr}):
                oneStockCloseList = []
                oneStockCloseList.append(stockIdStr)
                for i in range(1, totalSelectedDateNum):
                    try:
                        oneStockCloseList.append(stockDailyCloseResDict[dateList[i]])
                    except:
                        oneStockCloseList.append('0')
                totalStockCloseList.append(oneStockCloseList)


        #######################################
        # Retrieve indices close data
        #######################################
        stockListSize = len(totalStockCloseList)
        for stockIdRes in collection_indicesclose.find({}, {'StockID': 1}):
            stockIdStr = stockIdRes['StockID']
            for indexDailyCloseResDict in collection_indicesclose.find({'StockID': stockIdStr}):
                oneStockCloseList = []
                oneStockCloseList.append(stockIdStr)
                for i in range(1, totalSelectedDateNum):
                    try:
                        oneStockCloseList.append(indexDailyCloseResDict[dateList[i]])
                    except:
                        oneStockCloseList.append('0')
                totalStockCloseList.append(oneStockCloseList)

        print('Finished retrieving close data from MongoDB...')


        f = open(txtFileName, 'w')

        txtFileRowNum = len(oneStockCloseList)
        txtFileColNum = len(totalStockCloseList)

        f.write(str(txtFileRowNum - 1) + '\t' + str(txtFileColNum - 1) + '\t')

        for firstRowLen in range(txtFileColNum - 1):
            f.write('\t')
        f.write('\n')
        print('Finished writing text file header...')

        for rowIter in range(txtFileRowNum):
            for colIter in range(txtFileColNum):
                try:
                    cellData = totalStockCloseList[colIter][rowIter]
                    if math.isnan(float(cellData)):
                        cellData = 0
                    f.write(str(cellData) + '\t')
                except ValueError:
                    cellData = str(totalStockCloseList[colIter][rowIter])
                    f.write(cellData + '\t')

            f.write('\n')
            print('Wrote row number # ' + str(rowIter) + ' to the text file....')

        f.close()
        print('Finished writing the whole text file!')



    ###################################################################################################################
    ### Obtain Last Stock Update Date
    ###################################################################################################################
    def retrieveDataUpdateDate(self, strUpdateTimeKey):
        DBIP = self.DBIP
        DBPort = self.DBPort
        DBName = self.DBName
        colSysData = self.colSysData

        client = MongoClient(DBIP, DBPort)
        db = client[DBName]
        collection_sys_data = db[colSysData]


        sysDataRes = collection_sys_data.find({}, {strUpdateTimeKey: 1})

        dateListDict = sysDataRes[0]

        dataUpdateDate = dateListDict[strUpdateTimeKey]

        return (dataUpdateDate)


    ###################################################################################################################
    ### Update Date Data as Current Date
    ###################################################################################################################
    def updateDataUpdateDate(self, strUpdatedKeyName, strUpdatedKeyAlias, currentDateStr = time.strftime('%Y-%m-%d')):
        DBIP = self.DBIP
        DBPort = self.DBPort
        DBName = self.DBName
        colSysData = self.colSysData

        client = MongoClient(DBIP, DBPort)
        db = client[DBName]
        collection_sys_data = db[colSysData]

        try:
            collection_sys_data.update({}, {'$set': {strUpdatedKeyName: currentDateStr}})
        except:
            print('Failed to update stock update date to MongoDB, please check.......')
            return()
        print('Successfully updated ' + strUpdatedKeyAlias + ' into MongoDB.')





#######################################################################################################################
# 2008-1-1 is the starting date of the daily close data
#######################################################################################################################
obj = DailyClose(startDateStr='2008-01-01')



#######################################################################################################################
#Set the ending date for inserting and the starting date for updating
#Please do not run the following code unless database is corrupted.
#######################################################################################################################
baseEndDateStr = '2017-05-01'
#obj.writeBaseStocksFromWindToMongoDB(baseEndDateStr)
#obj.writeBaseIndicesFromWindToMongoDB(baseEndDateStr)

#######################################################################################################################
#The following two functions should be run before close data can be used for the backtrace application,
#######################################################################################################################

#obj.updateStocksFromWindToMongoDB(obj.retrieveDataUpdateDate('last_stock_daily_close_update_date'))
#obj.updateIndicesFromWindToMongoDB(obj.retrieveDataUpdateDate('last_index_daily_close_update_date'))



backtraceTxtFileStartDateStr = '2014-01-01'
obj.writeFromMongoDBtoTxt(backtraceTxtFileStartDateStr)
w.stop()
